{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'pandas'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[4], line 4\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;66;03m# 安装 Tushare 库（如果还没有安装）\u001b[39;00m\n\u001b[0;32m      2\u001b[0m \u001b[38;5;66;03m# !pip install tushare\u001b[39;00m\n\u001b[1;32m----> 4\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mpandas\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mpd\u001b[39;00m\n\u001b[0;32m      5\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mtf\u001b[39;00m\n\u001b[0;32m      6\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'pandas'"
     ]
    }
   ],
   "source": [
    "# 安装 Tushare 库（如果还没有安装）\n",
    "# !pip install tushare\n",
    "\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from numpy import array\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.layers import Dense, LSTM\n",
    "import mplfinance as mpf\n",
    "from mplfinance.original_flavor import candlestick2_ohlc\n",
    "from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score\n",
    "import tushare as ts  # 导入 Tushare 库\n",
    "\n",
    "# 设置随机种子\n",
    "np.random.seed(1)\n",
    "tf.random.set_seed(1)\n",
    "\n",
    "# 获取数据\n",
    "# 请确保你的API Key是有效的，并且不要在公共代码中暴露它\n",
    "pro = ts.pro_api('4963752acd3a2be9cafbabe76bacc8481e8c0ba9bcb692c5584846e2')\n",
    "df = pro.daily(ts_code='601939.SH', start_date='20100101', end_date='20241225')\n",
    "df.to_csv(\"stock601939.csv\")\n",
    "\n",
    "# 读取并处理数据\n",
    "mydata = pd.read_csv(\"stock601939.csv\", parse_dates=[\"trade_date\"], index_col=\"trade_date\")[[\"open\", \"high\", \"low\", \"close\"]]\n",
    "mydata = mydata.sort_index(ascending=True)  # 确保数据按时间顺序排列\n",
    "\n",
    "# 绘制K线图\n",
    "fig, ax = plt.subplots(figsize=(14, 7))\n",
    "candlestick2_ohlc(ax, opens=mydata['open'].values, highs=mydata['high'].values, lows=mydata['low'].values, closes=mydata['close'].values, width=0.5, colorup=\"r\", colordown=\"g\")\n",
    "\n",
    "# 添加均线\n",
    "mydata['5'] = mydata['close'].rolling(window=5).mean()\n",
    "mydata['10'] = mydata['close'].rolling(window=10).mean()\n",
    "\n",
    "plt.plot(mydata['5'].values, alpha=0.5, label='MA5')\n",
    "plt.plot(mydata['10'].values, alpha=0.5, label='MA10')\n",
    "\n",
    "ax.legend(facecolor='white', edgecolor='white', fontsize=10)\n",
    "plt.title('K线图及均线')\n",
    "plt.xlabel('日期')\n",
    "plt.ylabel('价格')\n",
    "plt.grid(True)\n",
    "plt.show()\n",
    "\n",
    "# 参数设置\n",
    "n_timestamp = 5\n",
    "n_epochs = 30\n",
    "\n",
    "# 数据准备\n",
    "training_set = mydata.iloc[:, 3:4].to_numpy()  # 使用 'close' 列\n",
    "test_set = mydata.iloc[-300:, 3:4].to_numpy()  # 使用最后300天作为测试集\n",
    "\n",
    "sc = MinMaxScaler(feature_range=(0, 1))\n",
    "training_set_scaled = sc.fit_transform(training_set)\n",
    "testing_set_scaled = sc.transform(test_set)\n",
    "\n",
    "# 数据切分\n",
    "def data_split(sequence, n_timestamp):\n",
    "    X = []\n",
    "    y = []\n",
    "    for i in range(len(sequence) - n_timestamp):\n",
    "        seq_x, seq_y = sequence[i:i+n_timestamp], sequence[i+n_timestamp]\n",
    "        X.append(seq_x)\n",
    "        y.append(seq_y)\n",
    "    return np.array(X), np.array(y)\n",
    "\n",
    "X_train, y_train = data_split(training_set_scaled, n_timestamp)\n",
    "X_train = X_train.reshape(X_train.shape[0], X_train.shape[1], 1)\n",
    "X_test, y_test = data_split(testing_set_scaled, n_timestamp)\n",
    "X_test = X_test.reshape(X_test.shape[0], X_test.shape[1], 1)\n",
    "\n",
    "# 模型构建\n",
    "model = Sequential()\n",
    "model.add(LSTM(units=50, activation='relu', input_shape=(X_train.shape[1], 1)))\n",
    "model.add(Dense(units=1))\n",
    "\n",
    "model.compile(optimizer=tf.keras.optimizers.Adam(0.001), loss='mean_squared_error')\n",
    "\n",
    "# 训练模型\n",
    "history = model.fit(X_train, y_train, batch_size=64, epochs=n_epochs, validation_data=(X_test, y_test), validation_freq=1)\n",
    "\n",
    "# 绘制损失曲线\n",
    "plt.figure(figsize=(14, 7))\n",
    "plt.plot(history.history['loss'], label='Training Loss')\n",
    "plt.plot(history.history['val_loss'], label='Validation Loss')\n",
    "plt.title('训练和验证损失')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()\n",
    "\n",
    "# 预测\n",
    "predicted_stock_price = model.predict(X_test)\n",
    "predicted_stock_price = sc.inverse_transform(predicted_stock_price)\n",
    "real_stock_price = sc.inverse_transform(y_test)\n",
    "\n",
    "# 绘制预测结果\n",
    "plt.figure(figsize=(14, 7))\n",
    "plt.plot(real_stock_price, color='red', label='实际股票价格')\n",
    "plt.plot(predicted_stock_price, color='blue', label='预测股票价格')\n",
    "plt.title('股票价格预测')\n",
    "plt.xlabel('时间')\n",
    "plt.ylabel('股票价格')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()\n",
    "\n",
    "# 计算评价指标\n",
    "MSE = mean_squared_error(real_stock_price, predicted_stock_price)\n",
    "RMSE = mean_squared_error(real_stock_price, predicted_stock_price) ** 0.5\n",
    "MAE = mean_absolute_error(real_stock_price, predicted_stock_price)\n",
    "R2 = r2_score(real_stock_price, predicted_stock_price)\n",
    "\n",
    "print('均方误差: %.5f' % MSE)\n",
    "print('均方根误差: %.5f' % RMSE)\n",
    "print('平均绝对误差: %.5f' % MAE)\n",
    "print('R2: %.5f' % R2)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
